#

#from InsightToolkit import *
import itk

from sys import argv
from math import pi

class reg():

    def __init__(self):
        """
        """
        self._fixedImageReader  = itk.ImageFileReader.IF2.New()
        self._movingImageReader = itk.ImageFileReader.IF2.New()

        #
        #  Instantiate the classes for the registration framework
        #
        self._registration    = itk.ImageRegistrationMethod.IF2IF2.New()
        #self._imageMetric     = itk.MeanSquaresImageToImageMetric.IF2IF2.New()
        #self._imageMetric     = itk.MattesMutualInformationImageToImageMetric.IF2IF2.New()
        self._imageMetric     = itk.NormalizedCorrelationImageToImageMetric.IF2IF2.New()
        #self._imageMetric     = itk.MeanReciprocalSquareDifferenceImageToImageMetric.IF2IF2.New()
        #self._imageMetric     = itk.MutualInformationImageToImageMetric.IF2IF2.New()
        self._transform       = itk.CenteredRigid2DTransform.D.New()
        self._optimizer       = itk.RegularStepGradientDescentOptimizer.New()
        self._interpolator    = itk.LinearInterpolateImageFunction.IF2D.New()

        self._registration.SetOptimizer(      self._optimizer.GetPointer() )
        self._registration.SetTransform(      self._transform.GetPointer() )
        self._registration.SetInterpolator(   self._interpolator.GetPointer() )
        self._registration.SetMetric(         self._imageMetric.GetPointer() )

        self._initialParameters = self._transform.GetParameters()
        self._initialParameters

        self._optimizer.SetNumberOfIterations( 1 )

        self._firstRun = True

    def SetFixedFileName(self,fixedImageFileName):
        """
        should only be required once (since the image is fixed, right)
        """
        self._fixedImageReader.SetFileName(  fixedImageFileName )
        self._fixedImageReader.Update()
        self._fixedImage  = self._fixedImageReader.GetOutput()
        self._registration.SetFixedImage(  self._fixedImage )
        self._registration.SetFixedImageRegion(  self._fixedImage.GetBufferedRegion() )

    def SetMovingFileName(self,movingImageFileName):
        """
        """
        self._movingImageFileName = movingImageFileName
        self._movingImageReader.SetFileName( self._movingImageFileName )


    def UpdateMovingImage(self):
        """
        Warning: This will only work if the moving image is consistently saved
        to the same filename!
        """
        self._movingImageReader.Update()
        self._movingImage = self._movingImageReader.GetOutput()

    def GetMetric(self):
        #
        # Read the fixed and moving images using filenames
        # from the command line arguments
        #

        self.UpdateMovingImage()
        self._registration.SetMovingImage( self._movingImage )


        self._initialParameters.fill(0)

        self._registration.SetInitialTransformParameters( self._initialParameters )
        #self._transform.SetParameters( self._initialParameters )

        #
        # Start the registration process
        #
##        startingParameters = self._transform.GetParameters()
##
##        print "Starting Registration Parameters "
##        print "Angle            =  %f" % (startingParameters.GetElement(0),)
##        print "Center X         =  %f" % (startingParameters.GetElement(1),)
##        print "Center Y         =  %f" % (startingParameters.GetElement(2),)
##        print "Translation X    =  %f" % (startingParameters.GetElement(3),)
##        print "Translation Y    =  %f" % (startingParameters.GetElement(4),)

        self._registration.StartRegistration()

        return self._optimizer.GetValue()

        #
        # Get the final parameters of the transformation
        #
        #finalParameters = registration.GetLastTransformParameters()


if __name__ == "__main__":
    fixedImageFileName = argv[1]
    movingImageFileName = argv[2]
    reg=reg()
    reg.GetMetric(fixedImageFileName, movingImageFileName)